from yoloSeg import YoloSegPose
import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
from pathlib import Path
import time

class PositionVisualizer:
    """
    用于可视化标准位置和测试位置的矩形
    
    使用YoloSegPose类的set_standard_position方法获取标准位置向量，
    使用test_one_image方法获取测试位置向量，
    返回向量结构为 - center: 转换后的中心点坐标
            - bottom_vec: 转换后的x方向向量
            - left_vec: 转换后的y方向向量
            - z_vec: 转换后的z方向向量
            - bottom_angle: bottom向量与x轴的夹角(度, 逆时针为正)
            - left_angle: left向量与y轴的夹角(度, 逆时针为正)
    以这些向量为左下角，绘制固定大小的矩形（长2.5米，宽0.8米）。
    """

    def __init__(self, seg_model_path, pose_model_path, standard_image_path, device=0, rect_size=(0.8, 2.5), scale_factor=100):
        """
        初始化PositionVisualizer类
        
        Args:
            seg_model_path: 分割模型路径
            pose_model_path: 姿态模型路径
            device: 设备ID
            rect_size: 矩形大小（宽, 长），单位为米，默认为(0.8, 2.5)
            scale_factor: 米到像素的转换因子，默认为100（1米=100像素）
        """
        self.yolo_seg_pose = YoloSegPose(seg_model_path, pose_model_path)
        self.standard_vector = self.yolo_seg_pose.set_standard_position(standard_image_path)
        self.rect_size = rect_size  # (宽度0.8米, 长度2.5米)
        self.scale_factor = scale_factor  # 米到像素的转换因子
        self.output_dir = "results_position_visualization"
        
        # 创建输出目录
        os.makedirs(self.output_dir, exist_ok=True)


    def calculate_position(self, image_path):
        """
        测试位置
        
        Args:
            image_path: 测试位置图像路径
            
        Returns:
            测试位置向量，包含以下结构：
            - center: 转换后的中心点坐标 (米)
            - bottom_vec: 转换后的x方向向量 (米)
            - left_vec: 转换后的y方向向量 (米)
            - z_vec: 转换后的z方向向量 (米)
            - bottom_angle: bottom向量与x轴的夹角(度, 逆时针为正)
            - left_angle: left向量与y轴的夹角(度, 逆时针为正)
        """
        pose,test_vector = self.yolo_seg_pose.test_one_image(image_path)
        if test_vector:
            return test_vector
        return None
    
    def draw_rectangle(self, vector, ax, label_prefix="Standard", color=(0, 1, 0), thickness=2):
        """
        使用Matplotlib在指定轴上绘制矩形、向量和角点，矩形固定为2.5米x0.8米
        
        Args:
            vector: 位置向量，包含以下结构：
                - center: 转换后的中心点坐标（矩形左下角点，米）
                - bottom_vec: 转换后的x方向向量（米）
                - left_vec: 转换后的y方向向量（米）
                - z_vec: 转换后的z方向向量（不使用）
                - bottom_angle: bottom向量与x轴的夹角（不使用）
                - left_angle: left向量与y轴的夹角（不使用）
            ax: Matplotlib轴对象
            label_prefix: 标签前缀（标准或测试）
            color: 矩形颜色 (RGB格式, 默认绿色，范围[0,1])
            thickness: 矩形线条粗细
        """
        if vector is None:
            return
        # 从向量中提取中心点坐标和向量（单位为米）
        center_x, center_y = vector['center'][0], vector['center'][1]
        left_vec = vector['left_vec'][:2]
        bottom_vec = vector['bottom_vec'][:2]
        
        # 归一化向量并缩放到固定尺寸（宽0.8米，长2.5米）
        rect_width, rect_height = self.rect_size  # 宽0.8米，长2.5米
        lx, ly = left_vec
        len_left = (lx**2 + ly**2)**0.5
        bx, by = bottom_vec
        len_bottom = (bx**2 + by**2)**0.5
        
        # 归一化并缩放left_vec到2.5米（长度）
        if len_left > 1e-5:
            left_vec = (lx / len_left * rect_height, ly / len_left * rect_height)
        else:
            left_vec = (0, rect_height)
        # 归一化并缩放bottom_vec到0.8米（宽度）
        if len_bottom > 1e-5:
            bottom_vec = (bx / len_bottom * rect_width, by / len_bottom * rect_width)
        else:
            bottom_vec = (rect_width, 0)
        
        # 应用像素缩放（1米 = scale_factor像素）
        center_x, center_y = center_x * self.scale_factor, center_y * self.scale_factor
        left_vec = (left_vec[0] * self.scale_factor, left_vec[1] * self.scale_factor)
        bottom_vec = (bottom_vec[0] * self.scale_factor, bottom_vec[1] * self.scale_factor)
        
        # 计算向量长度（像素）
        lx, ly = left_vec
        len_left = (lx**2 + ly**2)**0.5
        bx, by = bottom_vec
        len_bottom = (bx**2 + by**2)**0.5
        
        center = (center_x, center_y)  # 矩形左下角点（像素）
        # 计算与left_vec垂直的向量
        if len_left > 1e-5:
            perp1 = (ly / len_left, -lx / len_left)  # 顺时针旋转90°
            perp2 = (-ly / len_left, lx / len_left)  # 逆时针旋转90°
            dot1 = bx * perp1[0] + by * perp1[1]
            dot2 = bx * perp2[0] + by * perp2[1]
            if dot1 > dot2:
                adjusted_bottom_vec = (perp1[0] * len_bottom, perp1[1] * len_bottom)
            else:
                adjusted_bottom_vec = (perp2[0] * len_bottom, perp2[1] * len_bottom)
        else:
            adjusted_bottom_vec = bottom_vec

        # 计算矩形的四个角点（像素）
        bottom_left = center
        bottom_right = (center_x + adjusted_bottom_vec[0], center_y + adjusted_bottom_vec[1])
        top_left = (center_x + left_vec[0], center_y + left_vec[1])
        top_right = (top_left[0] + adjusted_bottom_vec[0], top_left[1] + adjusted_bottom_vec[1])
        
        # 调试：打印角点坐标
        print(f"{label_prefix}矩形角点: 左下={bottom_left}, 右下={bottom_right}, 左上={top_left}, 右上={top_right}")
        
        # 绘制矩形
        from matplotlib.patches import Polygon
        pts = np.array([bottom_left, bottom_right, top_right, top_left])
        polygon = Polygon(pts, closed=True, edgecolor=color, linewidth=thickness, fill=False, label=f"{label_prefix}Rect")
        ax.add_patch(polygon)
        
        # # 标记四个角点
        # ax.scatter(bottom_left[0], bottom_left[1], c='red', s=50, label=f"{label_prefix}左下")
        # ax.scatter(bottom_right[0], bottom_right[1], c='green', s=50, label=f"{label_prefix}右下")
        # ax.scatter(top_left[0], top_left[1], c='blue', s=50, label=f"{label_prefix}左上")
        # ax.scatter(top_right[0], top_right[1], c='yellow', s=50, label=f"{label_prefix}右上")
        
        # # 绘制left_vec（蓝色箭头）
        # ax.arrow(center_x, center_y, left_vec[0] * 0.5, left_vec[1] * 0.5, 
        #          color='blue', width=thickness/2, head_width=thickness*5, head_length=thickness*10, 
        #          label=f"{label_prefix}left_vec")
        
        # # 绘制adjusted_bottom_vec（红色箭头）
        # ax.arrow(center_x, center_y, adjusted_bottom_vec[0] * 0.5, adjusted_bottom_vec[1] * 0.5, 
        #          color='red', width=thickness/2, head_width=thickness*5, head_length=thickness*10, 
        #          label=f"{label_prefix}bottom_vec")
        
    # def draw_visualization(self, standard_vector, test_vector, min_canvas_size=(500, 500)):
    #     """
    #     绘制标准和测试位置的矩形、向量和比例尺
        
    #     Args:
    #         standard_vector: 标准位置向量 (米)
    #         test_vector: 测试位置向量 (米)
    #         min_canvas_size: 最小画布大小 (宽, 高)，默认为(500, 500)像素
    #     """
    #     # 计算画布大小
    #     rect_width, rect_height = self.rect_size
    #     max_vec_length = max(rect_width, rect_height) * self.scale_factor * 2
    #     canvas_width = max(min_canvas_size[0], int(max_vec_length + 100))
    #     canvas_height = max(min_canvas_size[1], int(max_vec_length + 150))
        
    #     # 创建Matplotlib图形
    #     fig, ax = plt.subplots(figsize=(canvas_width/100, canvas_height/100))
        
    #     # 绘制标准和测试矩形
    #     self.draw_rectangle(standard_vector, ax, label_prefix="标准", color=(0, 1, 0), thickness=2)
    #     self.draw_rectangle(test_vector, ax, label_prefix="测试", color=(0, 0, 1), thickness=2)
        
    #      # 添加比例尺文字
    #     scale_length_pixels = 100
    #     original_scale_length = scale_length_pixels / self.scale_factor
    #     scale_label = f"比例尺: {scale_length_pixels} 像素 = {original_scale_length:.2f} 米"
    #     print(f"{scale_label}")
        
        
    #     # 设置坐标轴范围（Y轴上正下负）
    #     ax.set_xlim(-100, canvas_width)
    #     ax.set_ylim(-100, canvas_height)
    #     ax.set_aspect('equal')
    #     ax.legend()
        
        
    #     return fig, ax


import os
import time
import shutil
import matplotlib.pyplot as plt
from datetime import datetime

def ensure_directory_exists(directory):
    """确保目录存在，如果不存在则创建"""
    if not os.path.exists(directory):
        os.makedirs(directory)

def move_to_processing(color_src, depth_src, filename):
    """将文件移动到处理目录"""
    processing_color_dir = "data/processing/Color"
    processing_depth_dir = "data/processing/Depth"
    
    ensure_directory_exists(processing_color_dir)
    ensure_directory_exists(processing_depth_dir)
    
    color_dst = os.path.join(processing_color_dir, filename)
    depth_dst = os.path.join(processing_depth_dir, filename)
    
    # 移动文件
    shutil.move(color_src, color_dst)
    if depth_src and os.path.exists(depth_src):
        shutil.move(depth_src, depth_dst)
    else:
        print(f"警告: 未找到对应的深度图 {depth_src}")
        depth_dst = None
    
    return color_dst, depth_dst

def cleanup_processing_files(color_path, depth_path):
    """清理处理后的文件"""
    if color_path and os.path.exists(color_path):
        os.remove(color_path)
    if depth_path and os.path.exists(depth_path):
        os.remove(depth_path)

def get_latest_image_pair(color_dir, depth_dir):
    """
    获取最新的彩色图和对应的深度图对
    返回: (color_path, depth_path, filename) 或 (None, None, None)
    """
    # 获取彩色图目录中所有PNG文件并按修改时间排序（最新的在前）
    color_files = [f for f in os.listdir(color_dir) if f.endswith('.png')]
    if not color_files:
        return None, None, None
    
    # 按修改时间排序
    color_files.sort(key=lambda x: os.path.getmtime(os.path.join(color_dir, x)), reverse=True)
    
    for filename in color_files:
        color_path = os.path.join(color_dir, filename)
        depth_path = os.path.join(depth_dir, filename)
        
        # 确保深度图存在
        if os.path.exists(depth_path):
            return color_path, depth_path, filename
    
    return None, None, None

def main():
    try:
        # 模型路径
        seg_model_path = "./merge_seg_epoch201_rknn_model"
        pose_model_path = "./merge_pose_epoch600_rknn_model"
        
        print(f"使用分割模型: {seg_model_path}")
        print(f"使用姿态模型: {pose_model_path}")

        # 标准位置图像路径
        standard_image_path = "data/Color1/440.png"
        print(f"标准位置图像路径: {standard_image_path}")
        visualizer = PositionVisualizer(seg_model_path, pose_model_path, standard_image_path, 
                                       rect_size=(0.8, 2.5), scale_factor=100)

        # 获取标准位置向量
        _,standard_vector = visualizer.calculate_position(standard_image_path)
        print(f"标准位置向量: {standard_vector}")

        # 设置Matplotlib交互模式
        plt.ion()  # 开启交互模式
        fig, ax = plt.subplots(figsize=(8, 8))  # 创建单一图形窗口
        plt.title("标准位置 (440.png) 与实时位置对比", fontsize=12)

        # 监控的图片目录
        color_dir = "data/Color3"
        depth_dir = "data/Depth3"
        processed_files = set()  # 用于记录已处理的文件
        
        # 设置坐标轴范围
        rect_width, rect_height = visualizer.rect_size
        max_vec_length = max(rect_width, rect_height) * visualizer.scale_factor 
        canvas_width = int(max_vec_length + 200)
        canvas_height = int(max_vec_length + 200)
        
        print(f"开始监控目录: {color_dir} 和 {depth_dir}")
        
        while True:
            try:
                # 获取最新的彩色图和深度图对
                color_src_path, depth_src_path, filename = get_latest_image_pair(color_dir, depth_dir)
                
                if not filename:
                    print("目录中没有有效的图像对，等待中...")
                    time.sleep(1)
                    continue
                
                # 如果文件已经处理过，则跳过
                if filename in processed_files:
                    time.sleep(0.1)  # 短暂休眠减少CPU占用
                    continue
                    
                print(f"\n处理最新图像对: {filename}")
                
                # 将文件移动到处理目录
                color_processing_path, depth_processing_path = move_to_processing(
                    color_src_path, depth_src_path, filename
                )
                
                # 获取测试位置向量
                pose,test_vector = visualizer.calculate_position(color_processing_path)
                processed_files.add(filename)
                cleanup_processing_files(color_processing_path, depth_processing_path)
                
                # 清空当前轴内容
                ax.clear()
                
                if test_vector is None:
                    print("无法计算位置向量，跳过此图像")
                    continue
                
                print(f"测试位置向量: {test_vector}")
                
                # 更新标题
                current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                ax.set_title(f"Standard (440.png) vs Test ({filename})\nUpdate Time: {current_time}", fontsize=12)
                
                # 绘制标准和测试矩形
                visualizer.draw_rectangle(standard_vector, ax, label_prefix="Standard", color=(0, 1, 0), thickness=2)
                visualizer.draw_rectangle(test_vector, ax, label_prefix="Test", color=(0, 0, 1), thickness=2)
                
                # 设置坐标轴范围
                ax.set_xlim(-50, canvas_width)
                ax.set_ylim(-50, canvas_height)
                ax.set_aspect('equal')
                ax.legend()
                
                # 刷新显示
                plt.draw()
                plt.pause(0.1)
                
                # 保存结果
                output_path = os.path.join(visualizer.output_dir, f"position_comparison_{filename}")
                fig.savefig(output_path, dpi=100, bbox_inches='tight')
                print(f"结果已保存至: {output_path}")
                
                # 将文件标记为已处理
                processed_files.add(filename)
                
                # 清理处理文件
                cleanup_processing_files(color_processing_path, depth_processing_path)
                
                # 限制已处理文件集合的大小，防止内存泄漏
                if len(processed_files) > 100:
                    # 按处理时间排序，移除最早处理的50个文件记录
                    oldest_files = sorted(
                        [(f, os.path.getmtime(os.path.join(color_dir, f))) for f in processed_files],
                        key=lambda x: x[1]
                    )[:50]
                    for f, _ in oldest_files:
                        processed_files.remove(f)
                        
            except KeyboardInterrupt:
                print("用户中断，退出监控...")
                break
            except Exception as e:
                print(f"处理过程中发生错误: {e}")
                time.sleep(1)  # 出错后稍作等待
                
    except Exception as e:
        import traceback
        print(f"执行过程中发生错误: {e}")
        print(traceback.format_exc())
        
    finally:
        # 确保图形窗口关闭
        plt.ioff()
        plt.close('all')
        print("程序结束")

if __name__ == "__main__":
    main()